import os
from dotenv import load_dotenv, find_dotenv # 导入 find_dotenv 帮助定位
from langchain.chains import ConversationChain
from langchain.memory import ConversationBufferMemory
from langchain.prompts import PromptTemplate
from langchain.llms import OpenAI

# 加载 .env 文件中的环境变量 (增强调试)
load_dotenv(dotenv_path=find_dotenv(usecwd=True), verbose=True, override=True)

# 从环境变量加载 API 密钥和基础 URL
api_key = os.getenv("OPENAI_API_KEY")
api_base = os.getenv("OPENAI_API_BASE")
os.environ["OPENAI_API_KEY"] = api_key
os.environ["OPENAI_API_BASE"] = api_base

# 自定义AI前缀
llm=OpenAI(
    temperature=0,
)

template = """下面是一段AI与人类的对话，AI会针对人类问题，提供尽可能详细的回答，如果AI不知道答案，会直接回复'人类老爷，我真的不知道'.
当前对话:
{history}
Human:{input}
AI助手:"""

prompt = PromptTemplate(
    template=template,
    input_variables=["history", "input"],
)

chain = ConversationChain(
    llm=llm,
    memory=ConversationBufferMemory(
        ai_prefix="AI助手",
        return_messages=True,
    ),
    prompt=prompt,
    verbose=True,
)

chain.predict(input="你知道我叫什么名字吗?")